Description of bug:

• Users can wrap their modules with helper classes like DataParallelCriterion or DataParallel

custom_loss_module = CustomLossModule()
parallel_custom_loss_module = DataParallelCriterion(custom_loss_module)

• The smdebug hook register each module like this:

        for name, submodule in module.named_modules():
            assert submodule not in self.module_set, f"Don't register module={module} twice"
            submodule._module_name = name
            self.module_set.add(submodule)
        module._module_name = module._get_name()
        self.module_set.add(module)

The problem with the above line is that the _module_name attribute is attached to the parallel_custom_loss_module and not the nested custom_loss_module.

• We would have instead needed to do:

submodule.module._module_name = name

• When the call to the forward hook is made, the helper module will internally call:

    def forward(self, *inputs, **kwargs):
        if not self.device_ids:
            return self.module(*inputs, **kwargs)

We should instead simply use:
module._get_name() in the forward_hook

Style and formatting:

I have run pre-commit install to ensure that auto-formatting happens with every commit.

Issue number, if available

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.